from phylopy.tree import Node, Tree
from pyparsing import *

"""Parse phylogenetic trees entered in Comapct Newick Format.  CNF is a subset of Newick
(New Hampshire) format that is more fully specified, and as a result, easier to parse without 
ambiguity."""

def preprocessString(cnfString, label):
    """Handles the case where a tree is provided which has missing terminal leaves, for instance
    ((,,),)
    Insert a set of dummy nodes for later processing steps."""
    C1 = Literal("(,")
    C2 = Literal(",)")
    C3 = Literal(",(")
    C4 = Literal("),")
    C5 = Literal(",,")
    grammar = (C1 | C2 | C3 | C4 | C5).setParseAction(lambda t: t[0][0] + label + t[0][1])
    new = grammar.transformString(cnfString)
    if new == cnfString:
        return new
    else:
        return preprocessString(new, label)


def parseTree(cnfString):
    label = Word(alphanums + "_")
    distance = (Suppress(":") + Word(nums + "+-.eE")).setParseAction(lambda t: float(t[0]))
    node = Forward()
    children = Suppress("(") + delimitedList(node) + Suppress(")")
    
    nodeC = children("children") + Optional(label)("label") + Optional(distance)("distance")
    nodeL = Optional(children)("children") + label("label") + Optional(distance)("distance")
    nodeD = Optional(children)("children") + Optional(label)("label") + distance("distance")
    node << Group(nodeC | nodeL | nodeD)
    tree = node + Optional(Suppress(";")) + stringEnd
    
    return tree.parseString(cnfString)


def processNodes(g, parent, res):
    """Recurses down the parse tree, creating Node objects and hooking them up to other graph
    elements with edges, as appropriate.  Note that recursion is probably safe here.  To reach
    Python's default recursion limit (99), one would have to have a tree with at least 
    exp(99 * ln(2)) leaves.  That's approaching the number of atoms in the universe, so if my
    calculations are correct, I think we are safe."""
    for elt in res:
        n = Node(elt.label)
        g.add_node(n)
        g.add_edge(parent, n, weight=elt.distance)
        if elt.children:
            processNodes(g, n, elt.children)


def makeTree(s):
    res = parseTree(s)
    g = Tree()
    res = res[0]
    parent = Node(res.label)
    g.add_node(parent)
    processNodes(g, parent, res.children)
    return g 


def runTests():
    """A cursory test harness."""
    tests = """\
        dog;
        (dog, cat, mouse)mammal;
        (dog, cat, (mouse, rat)rodent, gorilla)mammal;
        dog:10.1;
        ((dog:3, cat:7):4, mouse:8);""".splitlines()
    for t in tests:
        print t.strip() 
        t = makeTree(t)
        print t, [e.label for e in t.leaves()]


def loadTree(treeFile):
    f = open(treeFile, 'r')
    t = makeTree(f.read())
    f.close()
    return t

